import bpy
import os, shutil
import numpy as np
from argparse import ArgumentParser
parser = ArgumentParser(description='bpy arg parser')
parser.add_argument( "--mesh_type", type=str, required=True)
parser.add_argument("--mesh_list", nargs="+", type=str, default=[])
parser.add_argument('--debug_mode', type=int, default=10)
parser.add_argument('--debug_video_step', type=int, default=5)
parser.add_argument('--write_cover', action='store_true')
args = parser.parse_args()

mesh_type = args.mesh_type
mesh_list = args.mesh_list
cam_pose_dir = '/workspace/data/blender/cam_pose'
save_root_dir = '/workspace/data/blender/pc_res'
bpy.ops.wm.open_mainfile(filepath="/workspace/data/gs_data/pc.blend")
# bpy.data.scenes['Scene'].render.engine = 'BLENDER_WORKBENCH'

for mesh in mesh_list:
    bpy.data.objects[mesh].hide_render = False

    bpy.data.objects['Camera'].rotation_mode = 'QUATERNION'
    bpy.data.cameras['Camera'].lens_unit = 'FOV'
    
    # "bicycle".split('.')[0] = "bicycle", okay
    mesh_name = mesh.split('.')[0]
    extr_dir = os.path.join(cam_pose_dir, mesh_type, 'extrinsics')
    intr_dir = os.path.join(cam_pose_dir, mesh_type, 'intrinsics')
    save_dir = os.path.join(save_root_dir, mesh_type, mesh_name)

    if os.path.exists(save_dir) and not args.write_cover:
        assert False
    os.makedirs(save_dir, exist_ok=True)
    
    img_list = os.listdir(intr_dir)
    for img_id, img in enumerate(img_list):

        if args.debug_mode > 0 and img_id % args.debug_video_step != 0:
            continue
        img_name = img.split('.')[0]

        data_rot = np.loadtxt(os.path.join(extr_dir, img_name + '_q.txt'))
        data_loc = np.loadtxt(os.path.join(extr_dir, img_name + '_t.txt'))
        data_intr = np.loadtxt(os.path.join(intr_dir, img))
        data_name = img_name + '.png'
        
        for i in range(3):
            bpy.data.objects['Camera'].rotation_quaternion[i] = data_rot[i]
            bpy.data.objects['Camera'].location[i] = data_loc[i]
        bpy.data.objects['Camera'].rotation_quaternion[3] = data_rot[3]

        # bpy.data.cameras['Camera'].lens = data_intr[0]
        bpy.data.cameras['Camera'].angle = data_intr[0]
        if data_intr[1] > 1:
            bpy.data.scenes['Scene'].render.pixel_aspect_x = data_intr[1]
        else:
            bpy.data.scenes['Scene'].render.pixel_aspect_x = 1
            bpy.data.scenes['Scene'].render.pixel_aspect_y = 1 / data_intr[1]
        bpy.data.scenes['Scene'].render.pixel_aspect_y = 1

        downsample_rate1 = 1
        bpy.data.scenes['Scene'].render.resolution_x = int(data_intr[2]/downsample_rate1)
        bpy.data.scenes['Scene'].render.resolution_y = int(data_intr[3]/downsample_rate1)

        bpy.ops.render.render() # render an image
        bpy.data.images["Render Result"].save_render(filepath=os.path.join(save_dir, data_name)) # save the image

        if args.debug_mode > 0 and img_id > args.debug_mode * args.debug_video_step:
            break
    
    bpy.data.objects[mesh].hide_render = True
    # # clean up
    # if mesh_type.startswith('mip'):
    #     bpy.data.objects.remove(bpy.data.objects[mesh_name])
    
    # bpy.ops.outliner.orphans_purge(do_recursive=True)